{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "import os\n",
    "import sys\n",
    "import pickle\n",
    "import torch.nn.functional as F\n",
    "module_path = os.path.abspath(os.path.join('..'))\n",
    "if module_path not in sys.path:\n",
    "    sys.path.append(module_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model import EncoderRNN,DecoderRNN\n",
    "from config import *\n",
    "from pre_process import load_data, get_batch, pack_seqs\n",
    "USE_CUDA = torch.cuda.is_available()\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "BASE_DIR = '../data/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _sequence_mask(sequence_length, max_len=None):\n",
    "    if max_len is None:\n",
    "        max_len = sequence_length.data.max()\n",
    "    batch_size = sequence_length.size(0)\n",
    "    seq_range = torch.range(0, max_len - 1).long()\n",
    "    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)\n",
    "    seq_range_expand = Variable(seq_range_expand)\n",
    "    if sequence_length.is_cuda:\n",
    "        seq_range_expand = seq_range_expand.cuda()\n",
    "    seq_length_expand = (sequence_length.unsqueeze(1)\n",
    "                         .expand_as(seq_range_expand))\n",
    "    return seq_range_expand < seq_length_expand\n",
    "\n",
    "\n",
    "def compute_loss(logits, target, length):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        logits: A Variable containing a FloatTensor of size\n",
    "            (batch, max_len, num_classes) which contains the\n",
    "            unnormalized probability for each class.\n",
    "        target: A Variable containing a LongTensor of size\n",
    "            (batch, max_len) which contains the index of the true\n",
    "            class for each corresponding step.\n",
    "        length: A Variable containing a LongTensor of size (batch,)\n",
    "            which contains the length of each data in a batch.\n",
    "    Returns:\n",
    "        loss: An average loss value masked by the length.\n",
    "    \"\"\"\n",
    "\n",
    "    # logits_flat: (batch * max_len, num_classes)\n",
    "    logits_flat = logits.view(-1, logits.size(-1))\n",
    "    # log_probs_flat: (batch * max_len, num_classes)\n",
    "    log_probs_flat = functional.log_softmax(logits_flat)\n",
    "    # target_flat: (batch * max_len, 1)\n",
    "    target_flat = target.view(-1, 1)\n",
    "    # losses_flat: (batch * max_len, 1)\n",
    "    losses_flat = -torch.gather(log_probs_flat, dim=1, index=target_flat)\n",
    "    # losses: (batch, max_len)\n",
    "    losses = losses_flat.view(*target.size())\n",
    "    # mask: (batch, max_len)\n",
    "    mask = _sequence_mask(sequence_length=length, max_len=target.size(1))\n",
    "    losses = losses * mask.float()\n",
    "    loss = losses.sum() / length.float().sum()\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "module 'torch.functional' has no attribute 'log_softmax'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-32-b4b55f14b522>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     48\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     49\u001b[0m \u001b[0;31m#         loss = loss_function(tag_score, y[sent2ix_sort].view(-1))\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 50\u001b[0;31m         \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompute_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtag_score\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0msent2ix_sort\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mview\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx_lens\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     51\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     52\u001b[0m         \u001b[0mls\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-28-a33f731af6ef>\u001b[0m in \u001b[0;36mcompute_loss\u001b[0;34m(logits, target, length)\u001b[0m\n\u001b[1;32m     31\u001b[0m     \u001b[0mlogits_flat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlogits\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mview\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlogits\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     32\u001b[0m     \u001b[0;31m# log_probs_flat: (batch * max_len, num_classes)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 33\u001b[0;31m     \u001b[0mlog_probs_flat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfunctional\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlog_softmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlogits_flat\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     34\u001b[0m     \u001b[0;31m# target_flat: (batch * max_len, 1)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     35\u001b[0m     \u001b[0mtarget_flat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtarget\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mview\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mAttributeError\u001b[0m: module 'torch.functional' has no attribute 'log_softmax'"
     ]
    }
   ],
   "source": [
    "train_data = pickle.load(open(BASE_DIR + 'train_data_p.pkl', 'rb'))\n",
    "word2index = pickle.load(open(BASE_DIR + 'word_index.pkl', 'rb'))\n",
    "tag2index = pickle.load(open(BASE_DIR + 'tag2ix.pkl', 'rb'))\n",
    "\n",
    "encoder = EncoderRNN(len(word2index), EMBEDDING_DIM, ENCODER_DIM)\n",
    "decoder = DecoderRNN(len(tag2index), len(tag2index) //3, DECODER_DIM,10)\n",
    "\n",
    "if USE_CUDA:\n",
    "    encoder = encoder.cuda()\n",
    "    decoder = decoder.cuda()\n",
    "\n",
    "decoder.init_weights()\n",
    "\n",
    "weight_bias = []\n",
    "for i in range(162):\n",
    "    if i == 0:\n",
    "        weight_bias.append(1)\n",
    "    else:\n",
    "        weight_bias.append(2)\n",
    "\n",
    "weight_bias = torch.FloatTensor(weight_bias)\n",
    "\n",
    "loss_function = nn.CrossEntropyLoss(ignore_index=0)\n",
    "enc_optim = optim.Adam(encoder.parameters(), lr=0.01)\n",
    "dec_optim = optim.Adam(decoder.parameters(), lr=0.01)\n",
    "\n",
    "for step in range(2):\n",
    "    losses = []\n",
    "    for i, batch in enumerate(get_batch(BATCH_SIZE, train_data)):\n",
    "        sents, tags = zip(*batch)\n",
    "        x = torch.tensor(sents, dtype=torch.long)\n",
    "        y = torch.tensor(tags, dtype=torch.long)\n",
    "\n",
    "        x_lens = torch.tensor(pack_seqs(sents), dtype=torch.int)\n",
    "        x_lens_sort, sent2ix_sort = torch.sort(x_lens, descending=True)\n",
    "        sent2ix_unsort, _ = torch.sort(sent2ix_sort)\n",
    "\n",
    "        # y_lens_sort, tag2ix_sort = torch.sort(x_lens, descending=True)\n",
    "        # tag2ix_unsort, _ = torch.sort(tag2ix_sort)\n",
    "\n",
    "        encoder.zero_grad()\n",
    "        decoder.zero_grad()\n",
    "\n",
    "        enc_output, enc_hidden = encoder(x[sent2ix_sort], x_lens_sort)\n",
    "        start_decode = torch.LongTensor([[1]] * BATCH_SIZE)\n",
    "\n",
    "        tag_score = decoder(start_decode, enc_hidden, enc_output, x_lens_sort)\n",
    "\n",
    "#         loss = loss_function(tag_score, y[sent2ix_sort].view(-1))\n",
    "        loss = compute_loss(tag_score, y[sent2ix_sort].view(-1), x_lens)\n",
    "\n",
    "        ls = loss.data\n",
    "        loss.backward()\n",
    "        losses.append(ls)\n",
    "\n",
    "        torch.nn.utils.clip_grad_norm_(encoder.parameters(), 5.0)\n",
    "        torch.nn.utils.clip_grad_norm_(decoder.parameters(), 5.0)\n",
    "\n",
    "        enc_optim.step()\n",
    "        dec_optim.step()\n",
    "\n",
    "        print('step {}'.format(step))\n",
    "        print('loss {}'.format(np.mean(losses)))\n",
    "        break\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:miniconda3]",
   "language": "python",
   "name": "conda-env-miniconda3-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
